from random import shuffle
import torch
from model_official import LeNet
import torchvision
import torch.nn as nn
import torchvision.transforms as transforms
import torch.optim as optim
import torch.utils
import logging
# 设置日志
logging.basicConfig(filename='activation_functions.log', level=logging.INFO)

def main():
    # 使用gpu进行计算
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    transform = transforms.Compose(
        [transforms.ToTensor(),
        transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))]
    )


    """
    数据结构
    训练集:50,000 张图像，按照 10 个类别平均分配，每类包含 5,000 张。
    测试集:10,000 张图像，每类包含 1,000 张。
    数据格式：
    每张图像存储为 3D 张量，形状为 (3, 32, 32)，对应 3 个通道(RGB)和图像大小。
    标签是 0 到 9 的整数，表示所属类别。
    """

    train_set = torchvision.datasets.CIFAR10(root="./data",train=True,
                                            download=False,transform=transform)

    train_loader = torch.utils.data.DataLoader(train_set,batch_size=36,
                                            shuffle=True,num_workers=0)

    test_set = torchvision.datasets.CIFAR10(root="./data",train=True,
                                            download=False,transform=transform)
    test_loader = torch.utils.data.DataLoader(test_set,batch_size=5000,
                                            shuffle=True,num_workers=0)

    
    # 构建一个迭代器,用于迭代出测试数据,因为CIFAR10的测试数据只有10000张,并且batch_size设置了10000,
    # 那么直接就可以迭代一次就可以了,准备测试数据,用于后面使用
    dataiter = iter(test_loader)
    test_images, test_labels = next(dataiter)
    test_images = test_images.to(device)
    test_labels = test_labels.to(device)
    
    """
    测试代码:用于测试下载的数据
    """
    # classes = ('plane', 'car', 'bird', 'cat',
    #            'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
    # import matplotlib.pyplot as plt
    # import numpy as np
    # def imshow(img):
    #     img = img / 2 + 0.5     # unnormalize
    #     npimg = img.numpy()
    #     plt.imshow(np.transpose(npimg, (1, 2, 0)))
    #     plt.show()
    # # show images
    # imshow(torchvision.utils.make_grid(images))
    # # print labels
    # print(' '.join(f'{classes[labels[j]]:5s}' for j in range(batch_size)))

    """
    定义不同的激活函数，写入日志
    """
    activation_functions = {
        "ReLU": nn.ReLU(),
        "Sigmoid": nn.Sigmoid(),
        "Tanh": nn.Tanh()
    }

    for name, activation in activation_functions.items():
        print(f"Training with {name} activation function.")

        """
        数据进行训练
        """
        # 实例化模型
        net = LeNet(activation_fn=activation)  
        net = net.to(device)
        # 损失函数
        loss_function = nn.CrossEntropyLoss()  
        loss_function = loss_function.to(device) 
        # 迭代器  optim 是pytorch的模块,其中Adam是一个优化算法
        optimizer = optim.Adam(net.parameters(), lr=0.001) 
        #训练轮数
        for epoch in range(5):    
            running_loss = 0.0
            for step , data in enumerate(train_loader, start=0):
                # get the inputs; data is a list of [inputs,labels]
                inputs, labels = data
                inputs = inputs.to(device)
                labels = labels.to(device) 
                optimizer.zero_grad()
                outputs = net(inputs)
                loss = loss_function(outputs,labels)
                loss.backward()
                optimizer.step()
                running_loss += loss.item()
                if step % 500 == 499:
                    with torch.no_grad():
                        outputs = net(test_images)
                        predict_y = torch.max(outputs, dim=1)[1]
                        accuracy = torch.eq(predict_y,test_labels).sum().item() / test_labels.size(0)
                        # 记录到日志
                        logging.info('%s, epoch: %d, step: %5d, train_loss: %.3f, test_accuracy: %.3f' % 
                                     (name, epoch + 1, step + 1, running_loss / 500, accuracy))
                        print('[%d, %5d] train_loss: %.3f test_accuracy:%.3f' % 
                            (epoch+1,step+1,running_loss / 500, accuracy))
                        running_loss = 0.0
        print("Finished Training")
        save_path = f'./Lenet_model/Lenet_{name}.pth'
        torch.save(net.state_dict(), save_path)

if __name__ == '__main__':
    main()
